# -*- coding: utf-8 -*-

import tensorflow as tf
import tqdm
import data3.config as config
import load_image
import data3.preprocessing as preprocessing


def weight_variable(shape):
    initial = tf.truncated_normal(shape, stddev=0.1)
    return tf.Variable(initial)


def bias_variable(shape):
    initial = tf.constant(0.1, shape=shape)
    return tf.Variable(initial)


def conv2d(x, W):
    return tf.nn.conv2d(x, W, strides=[1, 1, 1, 1], padding='SAME')


def max_pool_2x2(x):
    return tf.nn.max_pool(x, ksize=[1, 2, 2, 1],
                          strides=[1, 2, 2, 1], padding='SAME')


def add_conv_layer(input_images, w_shape):
    W_conv = weight_variable(w_shape)
    b_conv = bias_variable([w_shape[3]])
    return tf.nn.relu(conv2d(input_images, W_conv) + b_conv)


input_images = tf.placeholder(tf.float32, shape=[None, config.char_image_width, config.char_image_height],
                              name='input_images')
input_labels = tf.placeholder(tf.float32, shape=[None, config.label_size], name='input_labels')
keep_prob = tf.placeholder('float', name='keep_prob')

input_images_flat = tf.reshape(input_images / 255, [-1, config.char_image_width, config.char_image_height, 1])
conv1 = add_conv_layer(input_images_flat, w_shape=[3, 3, 1, 32])
conv2 = add_conv_layer(conv1, w_shape=[3, 3, 32, 64])
conv3 = add_conv_layer(conv2, w_shape=[3, 3, 64, 32])

# 全连接层
W_fc1 = weight_variable([config.char_image_width * config.char_image_height * 32, 128])
b_fc1 = bias_variable([128])
conv_pool4_flat = tf.reshape(conv3, [-1, config.char_image_width * config.char_image_height * 32])
h_fc1 = tf.nn.relu(tf.matmul(conv_pool4_flat, W_fc1) + b_fc1)
# drop 防止过拟合
h_fc1_drop = tf.nn.dropout(h_fc1, keep_prob)
# 输出层
W_fc2 = weight_variable([128, config.label_size])
b_fc2 = bias_variable([config.label_size])

conv_res = tf.nn.softmax(tf.matmul(h_fc1_drop, W_fc2) + b_fc2, name='output_label')

cross_entropy = -tf.reduce_sum(input_labels * tf.log(conv_res + 1e-8))
train_step = tf.train.AdamOptimizer(1e-3).minimize(cross_entropy)
correct_prediction = tf.equal(tf.argmax(conv_res, 1), tf.argmax(input_labels, 1))

accuracy = tf.reduce_mean(tf.cast(correct_prediction, 'float'))
saver = tf.train.Saver(max_to_keep=1)


def train(generator_images_and_labels, train_times=10000):
    with tf.Session() as sess:
        sess.run(tf.global_variables_initializer())
        sess.graph.finalize()
        for i in tqdm.trange(train_times):
            batch = next(generator_images_and_labels)
            if i % 100 == 0:
                train_accuracy = sess.run(accuracy, feed_dict={
                    input_images: batch[0], input_labels: batch[1], keep_prob: 1.0
                })
                print('step %d, training accuracy %g' % (i, train_accuracy))
            sess.run(train_step, feed_dict={
                input_images: batch[0], input_labels: batch[1], keep_prob: 0.5
            })

        saver.save(sess, config.model_name)


if __name__ == '__main__':
    next_batch = load_image.generator_next_batch(config.yzm_path, config.suffix, preprocessing.mapping_line_parser,
                                                 preprocessing.preprocessing, preprocessing.char2one_hot)
    train(next_batch)